import numpy as np
import matplotlib.pyplot as plt

#Read in saved simulation results
exp = 'perfect foresight'
#exp = 'unanticipated' 

avg_60 = np.load('Saved_Data/'+exp+'_beta=60.npy')
avg_80 = np.load('Saved_Data/'+exp+'_beta=80.npy')
avgs = [avg_60, avg_80]

p_res = avg_60.shape[0]

#initialize subplots
fig, ax = plt.subplots(ncols=2, figsize=(14,6), sharex=True, sharey=True)

P1_range = np.ones((p_res, 2, 3))
P1_range[:,:, 0] =  4
P1_range[:,:, 1:] = 1
P1_range[:,0,-1] = np.linspace(0.25, 1.25, p_res)
p_e = P1_range[:,0,-1]

ax[0].set_title("60% discount rate")
ax[1].set_title("80% discount rate")

ax[0].set_xlabel("Initial Energy Price")
ax[1].set_xlabel("Initial Energy Price")
ax[0].set_ylabel("Energy Intensity Relative to New Entrant")

for i, avg_intensity in enumerate(avgs):
    ax[i].plot(p_e, avg_intensity[:,0]/avg_intensity[:,3], label='incumbents with fixed capital', color='tab:blue', ls='-.')
    ax[i].plot(p_e, avg_intensity[:,1]/avg_intensity[:,3], label='incumbents with adjustment costs', color='tab:orange', ls='--')
    ax[i].plot(p_e, avg_intensity[:,2]/avg_intensity[:,3], label='incumbents with flexible capital', color='tab:green', ls='-')
    ax[i].axhline(1, color='k', ls=':', label='new entrants')

ax[1].legend()
ax[0].set_ylim(0.90, 1.3)
ax[0].set_xlim(0.25,  1)
ax[0].set_xticks([0.4, 0.6, 0.8, 1, 1.2])

plt.show()
